package com.minimalist.basic.config.interceptor;

import com.github.benmanes.caffeine.cache.Cache;
import com.minimalist.basic.config.bucket.IpRateLimitConfig;
import com.minimalist.basic.utils.Sm4Util;
import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.ConsumptionProbe;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

/**
 * IP基础限流拦截器
 *
 * @author 11987
 */
@Component
@Slf4j
@AllArgsConstructor
public class IpBasedRateLimitInterceptor implements HandlerInterceptor {

    private final Cache<Object, Object> cache;
    private final IpRateLimitConfig ipRateLimitConfig;

    @Override
    public boolean preHandle(@NotNull HttpServletRequest request
            , @NotNull HttpServletResponse response, @NotNull Object handler) throws Exception {
        String ip = getClientIp(request);
        String logMessage =
                String.format("请求IP: %s 请求地址: %s 速率限制为: %d/%ds",
                        ip, request.getRequestURI(), ipRateLimitConfig.getCapacity()
                        , ipRateLimitConfig.getDuration().toSeconds());
        log.info(logMessage);
        Bucket bucket = (Bucket) cache.get(ip, _ -> newBucket());

        ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(1);
        if (probe.isConsumed()) {
            response.addHeader("X-Rate-Limit-Remaining-IP"
                    , String.valueOf(probe.getRemainingTokens()));
            return true;
        } else {
            long waitForRefill = probe.getNanosToWaitForRefill() / 1_000_000_000;
            response.addHeader("X-Rate-Limit-Retry-After-Seconds-IP"
                    , String.valueOf(waitForRefill));
            String message = "请求过于频繁，请等待" + waitForRefill + "秒后重试！";
            String encrypt = Sm4Util.encrypt(message);
            response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
            response.getWriter().write(encrypt);
            return false;
        }
    }

    private String getClientIp(HttpServletRequest request) {
        String header = request.getHeader("X-Forwarded-For");
        if (header == null) {
            return request.getRemoteAddr();
        }
        return header.split(",")[0];
    }

    private Bucket newBucket() {
        Bandwidth build = Bandwidth.builder().capacity(ipRateLimitConfig.getCapacity())
                .refillIntervally(ipRateLimitConfig.getRefillInterval()
                        , ipRateLimitConfig.getDuration()).build();
        return Bucket.builder()
                .addLimit(build)
                .build();
    }
}
